"""6.4.0 Fix #2224: Ajout d'une colonne antenne_id et follower_id dans company

Revision ID: 28a56ce17b2f
Revises: 2409eb7e97a4
Create Date: 2022-03-07 17:25:02.765369

"""

# revision identifiers, used by Alembic.
revision = "28a56ce17b2f"
down_revision = "4f714e529447"

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import mysql
from caerp.models.company import Company
from caerp.models.user.user import User


def update_database_structure():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("company", sa.Column("antenne_id", sa.Integer(), nullable=True))
    op.add_column("company", sa.Column("follower_id", sa.Integer(), nullable=True))
    op.create_foreign_key(
        op.f("fk_company_antenne_id"),
        "company",
        "antenne_option",
        ["antenne_id"],
        ["id"],
        ondelete="set null",
    )
    op.create_foreign_key(
        op.f("fk_company_follower_id"),
        "company",
        "accounts",
        ["follower_id"],
        ["id"],
        ondelete="set null",
    )


def migrate_datas():
    from alembic.context import get_bind
    from zope.sqlalchemy import mark_changed
    from caerp_base.models.base import DBSESSION

    session = DBSESSION()
    conn = get_bind()

    query = (
        session.query(Company)
        .outerjoin(User, Company.employees)
        .outerjoin(User.userdatas)
    )

    populate_company_antenne_id(query, session, conn)
    populate_company_follower_id(query, session, conn)

    mark_changed(session)
    session.flush()


def populate_company_antenne_id(query, session, conn):
    for company in query.all():
        for employee in company.employees:
            if employee.userdatas is None:
                continue

            if employee.userdatas.situation_antenne_id != None:
                # Set the antenne_id from the first employee not NULL
                conn.execute(
                    "update company set antenne_id={0} where company.id={1}".format(
                        employee.userdatas.situation_antenne_id, company.id
                    )
                )
                break


def populate_company_follower_id(query, session, conn):
    for company in query.all():
        for employee in company.employees:
            if employee.userdatas is None:
                continue

            if employee.userdatas.situation_follower_id != None:
                conn.execute(
                    "update company set follower_id={0} where company.id={1}".format(
                        employee.userdatas.situation_follower_id, company.id
                    )
                )
                break


def upgrade():
    update_database_structure()
    migrate_datas()


def downgrade():
    op.drop_constraint(op.f("fk_company_follower_id"), "company", type_="foreignkey")
    op.drop_constraint(op.f("fk_company_antenne_id"), "company", type_="foreignkey")
    op.drop_column("company", "follower_id")
    op.drop_column("company", "antenne_id")
